package main

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"crypto/hmac"
	cryptoRand "crypto/rand"
	"crypto/sha512"
	"encoding/base64"
	"encoding/binary"
	"encoding/gob"
	"encoding/json"
	"errors"
	"fmt"
	"github.com/google/uuid"
	"github.com/xuperchain/contract-sdk-go/code"
	"github.com/xuperchain/contract-sdk-go/driver"
	bswabe2 "github.com/zrynuaa/cpabe06/bswabe"
	bswabe1 "github.com/zrynuaa/cpabe06_client/bswabe"
	"hash"
	"io"
	"log"
	mathRand "math/rand"
	"time"
)

type ContCert struct {
	Contno_id   string  `json:"Contno_id"`
	Contno_name string  `json:"Contno_name"`
	Com_id      string  `json:"Com_id"`
	insured     string  `json:"insured"`
	prem        float32 `json:"prem"`
	filehash    string  `json:"filehash"`
}

type CRstruct struct {
	Contno_id string
	Com_id    string
	Desc      string
}

type CPstruct struct {
	Contno_id string
	Com_id    string
	insured   string
	Desc      string
}

func (c *ContCert) Initialize(ctx code.Context) code.Response {
	creator, ok := ctx.Args()["creator"]
	if !ok {
		return code.Errors("missing creator")
	}
	err := ctx.PutObject([]byte("creator"), creator)
	if err != nil {
		return code.Error(err)
	}
	return code.OK(nil)
}

// 发布产品信息，这里不加密，只是存储到区块链
func ProductiontoTransaction(ctx code.Context) code.Response {

	initiator := ctx.Initiator()
	if initiator == "" {
		return code.Error(code.ErrMissingInitiator)
	}

	args := struct {
		Production_id string `json:"Production_id" validate:"required"`
		compamny_id   string `json:"compamny_id" validate:"required"`
		Desc          string `json:"desc" validate:"required"`
	}{}
	if err := code.Unmarshal(ctx.Args(), &args); err != nil {
		return code.Error(err)
	}

	ProductionKey := args.Production_id

	if err := ctx.PutObject([]byte(ProductionKey), []byte(args.Desc)); err != nil {
		return code.Error(err)
	}

	return code.OK([]byte(args.Production_id))

}

// 获取发布的产品信息，mark作为扩展字段，可以区分不同角色
func GetreadyInfo(ctx code.Context) code.Response {

	args := struct {
		mark          string `json:"make" validate:"required"`
		Production_id string `json:"Production_id" validate:"required"`
	}{}
	if err := code.Unmarshal(ctx.Args(), &args); err != nil {
		return code.Error(err)
	}
	mark := args.mark

	if mark == "RIN" {
		ProductionKey := args.Production_id
		key := []byte(ProductionKey)
		value, err := ctx.GetObject(key)
		if err != nil {
			log.Print("Production not found")
			return code.Error(err)
		}
		Production, _ := json.Marshal(value)
		return code.OK(Production)
	}
	return code.Error(errors.New("have no premission"))
}

// 加密并上传CR合同
func (e *ContCert) EncryptCR(ctx code.Context) code.Response {

	//01上传CR信息 获取前端传递的参数
	args := ctx.Args()

	//02 生成KDF派生需要的密钥
	//Mkey for KDF
	keyLength := 16
	MKey, err := generateRandomKey(keyLength)
	if err != nil {
		fmt.Println("Error generating random key:", err)
		return code.Error(err)
	}

	//03 将CR和密钥一起属性加密
	// ConPolicy := "A B C 1of3"
	//RWPolicy := "R 1of1 W R 1of2"
	FullPolicy := "A B C 1of3 R 1of1 W R 1of2 2of2" //test data
	var attrs = "A"
	var pub, _ = GetABEKeys(attrs) //获取ABE服务上的密钥对

	// transform to bytes
	var result bytes.Buffer
	temp := args
	encoder := gob.NewEncoder(&result)
	err1 := encoder.Encode(temp)
	if err1 != nil {
		log.Panic(err1)
	}
	result.Write(MKey)

	EnCRdata := bswabe1.SerializeBswabeCphKey(bswabe1.CP_Enc(pub, string(result.Bytes()), FullPolicy))
	//EnCRdata := bswabe1.SerializeBswabeCphKey(bswabe1.CP_Enc(pub,Desc,FullPolicy))

	//ContnoKey := "001"
	ContnoKey, ok := args["Contno"]
	if ok {
		if err := ctx.PutObject([]byte(ContnoKey), []byte(EnCRdata)); err != nil {
			fmt.Println("Error generating random key:", err)
			return code.Error(err)
		}
	} else {
		fmt.Println("Contno is null ", err)
	}

	//04 end
	return code.OK([]byte(ContnoKey))

}

// 获取CR合同并解密
func (e *ContCert) GetCR(ctx code.Context, Contno_id string, user string) code.Response {

	fmt.Println("Hello, World!")

	//01 获取密钥
	var attrs = "A"
	var pub, prv = GetABEKeys(attrs) //获取ABE服务上的密钥对
	// Init("001",pub,prv)

	//02 上传查询索引
	// args := struct {
	// 	Contno_id 		string    `json:"Contno_id" validate:"required"`
	// 	user 	 		string    `json:"user" validate:"required"`
	// }{}
	// if err := code.Unmarshal(ctx.Args(), &args); err != nil {
	// 	return code.Error(err)
	// }

	//03 获取数据
	CRContnoKey := Contno_id
	key := []byte(CRContnoKey)
	value, err := ctx.GetObject(key)
	if err != nil {
		fmt.Println("CR not found", err)
		return code.Error(err)
	}
	CR, _ := json.Marshal(value)

	//04 解密CR
	//result := bswabe.CP_Dec(pub, prv, CR)
	result := bswabe1.CP_Dec(pub, prv, bswabe1.UnSerializeBswabeCphKey(pub, []byte(CR)))

	//05 end
	return code.OK(result)
}

// 加密个人合同
func (e *ContCert) EncryptCP(ctx code.Context, Contno_id string, Com_id string, insured string, Desc string, CR_id string, user string) code.Response {

	// 01个人合同数据，需上传
	// args1 := struct {
	// 	Contno_id 		string    `json:"Contno_id"`
	// 	Com_id      	string    `json:"Com_id"`
	// 	insured  		string    `json:"insured"`
	// 	Desc 			string 	  `json:"desc" validate:"required"`
	// }{}
	// if err := code.Unmarshal(ctx.Args(), &args1); err != nil {
	// 	return code.Error(err)
	// }

	//02 解密CR合同，获取Mkey，生成对称密钥
	var attrs = "A"
	var pub, prv = GetABEKeys(attrs) //获取ABE服务上的密钥对
	// Init("001",pub,prv)//用户为001

	// args := struct {
	// 	CR_id 		    string    `json:"Contno_id" validate:"required"`
	// 	user 	 		string    `json:"user" validate:"required"`
	// }{}
	// if err := code.Unmarshal(ctx.Args(), &args); err != nil {
	// 	return code.Error(err)
	// }

	CRContnoKey := CR_id
	key := []byte(CRContnoKey)
	value, err := ctx.GetObject(key)
	if err != nil {
		return code.Error(err)
	}

	CR, _ := json.Marshal(value)

	//cphKey := (*bswabe.BswabeCphKey)(unsafe.Pointer(&CR))
	result := bswabe1.CP_Dec(pub, prv, bswabe1.UnSerializeBswabeCphKey(pub, []byte(CR)))
	//result := bswabe.CP_Dec(pk, sk, bswabe.UnSerializeBswabeCphKey(pk, []byte(j)))

	if string(result) != "" {

		if len(result) >= 16 {
			FullPolicy := "A B C 1of3 R 1of1 W R 1of2 2of2" //test data
			MKey := result[len(result)-16:]                 //获取后16位的MKey

			// 生成随机数作为salt
			salt := make([]byte, 16)
			_, err := mathRand.Read(salt)
			if err != nil {
				//fmt.Println("Error generating random bytes:", err)
				return code.Error(err)
			}

			var info []byte
			hkdf := NewHkdf(salt, MKey, sha512.New)
			AESkey := hkdf.Expand(info, 32, sha512.New) //最终的AES对称密钥

			UUID, err := uuid.NewRandom()
			UUID_ := UUID.String() + string(salt)

			//03 CPABE加密个人合同的salt
			EnCRdata_salt := bswabe1.SerializeBswabeCphKey(bswabe1.CP_Enc(pub, UUID_, FullPolicy))
			//再次使用CPABE加密

			ContnoKey_salt_key := Contno_id
			if err := ctx.PutObject([]byte(ContnoKey_salt_key), []byte(EnCRdata_salt)); err != nil {
				return code.Error(err)
			}

			var result bytes.Buffer
			var temp = CPstruct{
				Contno_id: Contno_id,
				Com_id:    Com_id,
				insured:   insured,
				Desc:      Desc,
			}

			//04 AES加密个人合同

			// temp := args1
			//    temp.Contno_id=Contno_id
			//    temp.Com_id=Com_id
			//    temp.insured=insured
			//    temp.Desc=Desc

			encoder := gob.NewEncoder(&result)
			err1 := encoder.Encode(temp)
			if err1 != nil {
				log.Panic(err1)
			}
			byteSlice := result.Bytes()
			ciphertext, err := encrypt([]byte(byteSlice), AESkey) //AES加密
			if err != nil {
				return code.Error(err)
			}

			// 将加密后的数据转换为Base64字符串以便传输或存储
			encoded := base64.StdEncoding.EncodeToString(ciphertext)

			Contno_id_key := Contno_id
			if err := ctx.PutObject([]byte(Contno_id_key), []byte(encoded)); err != nil {
				return code.Error(err)
			}

		} else {
			return code.Error(errors.New("something is wrong in encryption"))
		}
	} else {
		return code.Error(errors.New("You have no permission..."))
	}
	//05 end
	return code.OK([]byte("ok"))
}

// 解密个人合同
func (e *ContCert) DecryptCP(ctx code.Context) code.Response {
	//00 获取ABE服务上的密钥对
	var attrs = "A"
	var pub, prv = GetABEKeys(attrs)
	// Init("001",pub,prv)

	//01 解密CR合同，获取Mkey
	args := struct {
		CP_id string `json:"Contno_id" validate:"required"`
		user  string `json:"user" validate:"required"`
	}{}
	if err := code.Unmarshal(ctx.Args(), &args); err != nil {
		return code.Error(err)
	}

	CRContnoKey := args.CP_id
	key := []byte(CRContnoKey)
	value, err := ctx.GetObject(key)
	if err != nil {
		return code.Error(errors.New("CR not found"))
	}
	CR, _ := json.Marshal(value)

	//result := bswabe.CP_Dec(pub, prv, CR)
	result := bswabe1.CP_Dec(pub, prv, bswabe1.UnSerializeBswabeCphKey(pub, []byte(CR)))
	if string(result) != "" {

		if len(result) >= 16 {
			MKey := result[len(result)-16:] //获取后16位的MKey

			//02 获取个人合同的salt
			args2 := struct {
				uuid string `json:"uuid" validate:"required"`
				user string `json:"user" validate:"required"`
			}{}
			if err := code.Unmarshal(ctx.Args(), &args2); err != nil {
				return code.Error(err)
			}

			cp1 := args2.uuid
			key := []byte(cp1)
			value, err := ctx.GetObject(key)
			if err != nil {
				return code.Error(errors.New("cp1 not found"))
			}
			CP_salt, _ := json.Marshal(value)

			//03 生成对称密钥
			var info []byte
			hkdf := NewHkdf(CP_salt, MKey, sha512.New)
			AESkey := hkdf.Expand(info, 32, sha512.New) //最终的AES对称密钥

			//04 查询个人合同
			args1 := struct {
				Contno_id string `json:"Contno_id" validate:"required"`
				user      string `json:"user" validate:"required"`
			}{}
			if err := code.Unmarshal(ctx.Args(), &args1); err != nil {
				return code.Error(err)
			}

			CpContnoKey := args1.Contno_id
			key1 := []byte(CpContnoKey)
			value1, err := ctx.GetObject(key1)
			if err != nil {
				return code.Error(errors.New("CR not found"))
			}
			CP, _ := json.Marshal(value1)

			//05 解密个人合同

			CPdecoded, err := base64.StdEncoding.DecodeString(string(CP))
			if err != nil {
				fmt.Println("CP Decoding error:", err)
				return code.Error(err)
			}

			CP_mw, err := decrypt(CPdecoded, AESkey)
			if err != nil {
				fmt.Println("CP_mw Decryption error:", err)
				return code.Error(err)
			}
			return code.OK([]byte(CP_mw))

		} else {
			return code.Error(errors.New("something is wrong in CPencryption"))
		}
	} else {
		return code.Error(errors.New("You have no permission..."))
	}
	//06 end
	return code.OK([]byte("###"))
}

// 当合同到期后，撤销写权限
func (e *ContCert) Re_encrypt(ctx code.Context) code.Response {
	//01 获取并解析合同失效时间，判断是否失效
	compareDateStr := "2024-01-02"
	currentDate := time.Now()
	compareDate, err := time.Parse("2006-01-02", compareDateStr)
	if err != nil {
		return code.Error(err)
	}

	if compareDate.Before(currentDate) {

		//02 变更policy,删除RW树的左子树，这里直接输入了，可以代码实现删除左子树
		NewPolicy := "A B C 1of3 R 1of1 2of2"
		var attrs = "A"
		var pub, _ = GetABEKeys(attrs)
		//03 获取salt
		args2 := struct {
			uuid string `json:"uuid" validate:"required"`
			user string `json:"user" validate:"required"`
		}{}
		if err := code.Unmarshal(ctx.Args(), &args2); err != nil {
			return code.Error(err)
		}

		cp1 := args2.uuid
		key := []byte(cp1)
		value, err := ctx.GetObject(key)
		if err != nil {
			return code.Error(errors.New("cp1 not found"))
		}
		CP_salt, _ := json.Marshal(value)

		//04 重新加密CP_salt
		EnCRdata_salt := bswabe1.SerializeBswabeCphKey(bswabe1.CP_Enc(pub, string(CP_salt), NewPolicy))

		//05 上传到区块链
		ReEncrypt_salt_key := args2.uuid
		if err := ctx.PutObject([]byte(ReEncrypt_salt_key), []byte(EnCRdata_salt)); err != nil {
			return code.Error(err)
		}
	}
	//06 end
	return code.OK([]byte("ok"))
}

func generateRandomKey(length int) ([]byte, error) {
	key := make([]byte, length)

	_, err := mathRand.Read(key)
	if err != nil {
		return nil, err
	}

	return key, nil
}

// 获取ABE服务器上的主公钥
func GetABEKeys(attrs string) (*bswabe1.BswabePub, *bswabe1.BswabePrv) {

	//    client, err := rpc.DialHTTP("tcp", "127.0.0.1")
	//    if err != nil {
	// 	   log.Fatal("dialing:", err)
	//    }

	//    // Synchronous call同步方式调用
	//    var reply []byte
	//    err = client.Call("CPABE.Getpub", "", &reply)
	//    if err != nil {
	// 	   log.Fatal(" error:", err)
	//    }
	//    var pub = bswabe.UnSerializeBswabePub(reply)//获取PublicKey

	//    err = client.Call("CPABE.Getsk", attrs, &reply)
	//    if err != nil {
	// 	   log.Fatal("arith error:", err)
	//    }
	//    var prv = bswabe.UnSerializeBswabePrv(pub, reply) //获取服务端返回的解密私钥

	pub := new(bswabe2.BswabePub)
	msk := new(bswabe2.BswabeMsk)

	data := bswabe2.SerializeBswabePub(pub)
	ppuu := bswabe1.UnSerializeBswabePub(data) //获得服务端返回的公共参数

	prv1 := bswabe2.CP_Keygen(pub, msk, attrs)
	data = bswabe2.SerializeBswabePrv(prv1)
	pprr := bswabe1.UnSerializeBswabePrv(ppuu, data) //获取服务端返回的解密私钥

	return ppuu, pprr
}

// 用户初始化
//
//	func Init(name string, pub *bswabe.BswabePub, prv *bswabe.BswabePrv) error{
//		if name == "" {
//			return fmt.Errorf("name is nothing")
//		}
//		Name := name
//		sk := prv
//		pk := pub
//		return nil
//	}
//
// HKDF
func hkdfExtract(salt, inputKeyMaterial []byte, hashFunc func() hash.Hash) []byte {
	hashLen := hashFunc().Size()

	if salt == nil || len(salt) == 0 {
		salt = make([]byte, hashLen)
	}

	h := hmac.New(hashFunc, salt)
	h.Write(inputKeyMaterial)
	prk := h.Sum(nil)

	return prk
}

func hkdfExpand(pseudoRandomKey, info []byte, length int, hashFunc func() hash.Hash) []byte {
	hashLen := hashFunc().Size()
	length = length

	// 计算需要的块数
	blocksNeeded := (length + hashLen - 1) / hashLen
	okm := make([]byte, 0)
	outputBlock := make([]byte, 0)

	for counter := 0; counter < blocksNeeded; counter++ {
		// 计算每个块的输出
		h := hmac.New(hashFunc, pseudoRandomKey)
		h.Write(outputBlock)
		h.Write(info)
		counterBytes := make([]byte, 4)
		binary.BigEndian.PutUint32(counterBytes, uint32(counter+1))
		h.Write(counterBytes)
		outputBlock = h.Sum(nil)

		// 将输出块追加到最终输出
		okm = append(okm, outputBlock...)
	}

	return okm[:length]
}

type Hkdf struct {
	prk []byte
}

func NewHkdf(salt, inputKeyMaterial []byte, hashFunc func() hash.Hash) *Hkdf {
	prk := hkdfExtract(salt, inputKeyMaterial, hashFunc)

	return &Hkdf{prk: prk}
}

func (h *Hkdf) Expand(info []byte, length int, hashFunc func() hash.Hash) []byte {
	return hkdfExpand(h.prk, info, length, hashFunc)
}

// AES加密
func encrypt(plaintext []byte, key []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}

	// 创建GCM（Galois/Counter Mode）模式
	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}

	// 生成随机的nonce（一次性使用的值）
	nonce := make([]byte, gcm.NonceSize())
	if _, err := io.ReadFull(cryptoRand.Reader, nonce); err != nil {
		return nil, err
	}

	ciphertext := gcm.Seal(nil, nonce, plaintext, nil)
	result := append(nonce, ciphertext...)

	return result, nil
}

// AES解密
func decrypt(ciphertext []byte, key []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}

	gcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}

	nonceSize := gcm.NonceSize()
	if len(ciphertext) < nonceSize {
		return nil, fmt.Errorf("ciphertext is too short")
	}
	nonce, ciphertext := ciphertext[:nonceSize], ciphertext[nonceSize:]

	plaintext, err := gcm.Open(nil, nonce, ciphertext, nil)
	if err != nil {
		return nil, err
	}

	return plaintext, nil
}

func generateUUID() string {
	newUUID, err := uuid.NewRandom()
	if err != nil {
		fmt.Println("Error generating UUID:", err)
		return ""
	}
	return newUUID.String()
}

func main() {
	//  driver.Serve(newContCert())
	driver.Serve(new(ContCert))
	// fmt.Println("Hello from main function!")
}
